#pragma once
#include <atomic>
#include <cassert>
#include <condition_variable>
#include <functional>
#include <list>

#include "../refbase/refbase.h"

template<class T>
class Promise : public RefBase {
public:
    using PromiseT = Promise<T>;
    using sptrPromiseT = sptr<PromiseT>;
    using PromiseCompleteFunc = std::function<void(const T&)>;
    using PromiseRaceFunc = std::function<void(const PromiseT *)>;

    Promise() = default;
    virtual ~Promise() = default;
    Promise(const T& t);

    bool IsResolved() const;
    const T &Await();
    void Then(const PromiseCompleteFunc &func);
    void Resolve(const T &t);
    static void All(const std::list<sptrPromiseT> &promises);
    static const sptr<Promise<T>> &Race(const std::list<sptrPromiseT> &promises);

private:
    void SetRaceFunc(const PromiseRaceFunc &func);

    std::atomic<bool> resolved = false;
    std::mutex mutex;
    std::condition_variable cv;
    T value;
    PromiseCompleteFunc onComplete = nullptr;
    PromiseRaceFunc onRace = nullptr;
};

template<class T>
Promise<T>::Promise(const T& t)
    : value(t), resolved(true)
{
}

template<class T>
bool Promise<T>::IsResolved() const
{
    return resolved;
}

template<class T>
const T &Promise<T>::Await()
{
    if (resolved == false) {
        std::unique_lock<std::mutex> lock(mutex);
        cv.wait(lock, [this]() { return resolved == true; });
    }
    return value;
}

template<class T>
void Promise<T>::Then(const PromiseCompleteFunc &func)
{
    if (resolved == false) {
        std::unique_lock<std::mutex> lock(mutex);
        if (resolved == false) {
            onComplete = func;
        }
    } else {
        func(value);
    }
}

template<class T>
void Promise<T>::Resolve(const T &t)
{
    if (resolved == false) {
        std::unique_lock<std::mutex> lock(mutex);
        if (resolved == false) {
            value = t;
            resolved = true;

            if (onRace != nullptr) {
                onRace(this);
            }
            if (onComplete != nullptr) {
                onComplete(value);
            }
            cv.notify_all();
        }
    }
}

template<class T>
void Promise<T>::All(const std::list<sptrPromiseT> &promises)
{
    for (const auto &promise : promises) {
        promise->Await();
    }
}

template<class T>
const sptr<Promise<T>> &Promise<T>::Race(const std::list<sptrPromiseT> &promises)
{
    sptr<Promise<const PromiseT *>> racePromise = new Promise<const PromiseT *>();
    auto raceFunc = [racePromise](const PromiseT *that) { racePromise->Resolve(that); };
    for (const auto &promise : promises) {
        promise->SetRaceFunc(raceFunc);
    }

    const PromiseT *ret = racePromise->Await();
    for (const auto &promise : promises) {
        if (ret == promise.GetRefPtr()) {
            return promise;
        }
    }

    assert(0);
    return promises.front();
}

template<class T>
void Promise<T>::SetRaceFunc(const PromiseRaceFunc &func)
{
    if (resolved == false) {
        std::unique_lock<std::mutex> lock(mutex);
        if (resolved == false) {
            onRace = func;
        }
    } else {
        func(this);
    }
}
